import math
import os
import cv2
import numpy as np
import os.path as osp
import argparse
import csv
from pathlib import Path
import json
from tqdm import tqdm

# line_gt_path = r'/home/ych6szh/Documents/parking_structure/slot_train_0212/lines_train_new_0228.json'
box_gt_path = r'cwh_train_local.json'
save_dir = r'cwh_train_local_gt_view'
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

# with open(str(line_gt_path), 'r') as f_point:
#     line_result = json.load(f_point)

with open(str(box_gt_path), 'r') as f_box:
    box_result = json.load(f_box)

img_info = box_result['images']
# line_ann = line_result['annotations']
box_ann = box_result['annotations']
line_ann_new = []
box_ann_new = []

img_list = []

for single in tqdm(img_info):
    img_name = single['file_name']
    img_name = os.path.join('./cwh_train_local/',img_name)
    img_id = single['id']
    img = cv2.imread(img_name)
    if img is None:
        continue
    name = Path(img_name).name
    bbox_center_list = []
    # for single_line in line_ann:
    #     if single_line['image_id'] == img_id:
    #         line = single_line['line']
    #         cv2.line(img, (int(line[0][0]), int(line[0][1])), (int(line[1][0]), int(line[1][1])), (255, 0, 0), 3)
    for single_box in box_ann:
        if single_box['image_id'] == img_id:
            bbox = single_box['bbox']
            xc, yc, w, h, ag = bbox[:5]
            wx, wy = w / 2 * np.cos(ag), w / 2 * np.sin(ag)
            hx, hy = -h / 2 * np.sin(ag), h / 2 * np.cos(ag)
            p1 = [int(xc - wx - hx), int(yc - wy - hy)]
            p2 = [int(xc + wx - hx), int(yc + wy - hy)]
            p3 = [int(xc + wx + hx), int(yc + wy + hy)]
            p4 = [int(xc - wx + hx), int(yc - wy + hy)]
            cv2.line(img, tuple(p1), tuple(p2), (0, 255, 0), 2)
            cv2.line(img, tuple(p2), tuple(p3), (0, 255, 0), 2)
            cv2.line(img, tuple(p3), tuple(p4), (0, 255, 0), 2)
            cv2.line(img, tuple(p4), tuple(p1), (0, 255, 0), 2)     

    save_path = os.path.join(save_dir, name)
    cv2.imwrite(save_path, img)
#
# struct_json_point = {
#     'images': img_info,
#     'annotations': point_ann_new,
#     'categories': point_result['categories'],
# }

# result_file_path = Path(save_dir).joinpath('point_test_0225.json')
# json.dump(struct_json_point, open(str(result_file_path), "w"), indent=4)
#
# struct_json_box = {
#     'images': img_info,
#     'annotations': box_ann_new,
#     'categories': box_result['categories'],
# }
#
# result_file_path = Path(save_dir).joinpath('box_test_0225.json')
# json.dump(struct_json_box, open(str(result_file_path), "w"), indent=4)
# file.close()


